package com.zixuan.sina.ai;

import ai.onnxruntime.*;

import java.nio.FloatBuffer;
import java.nio.file.Paths;
import java.util.Collections;
import java.util.Map;

public class ONNXTextClassifier {
    private OrtSession session;

    public ONNXTextClassifier(String modelPath) throws OrtException {
        OrtEnvironment env = OrtEnvironment.getEnvironment();
        OrtSession.SessionOptions opts = new OrtSession.SessionOptions();
        session = env.createSession(Paths.get(modelPath).toString(), opts);
    }

    public float classifyText(float[] inputVector) throws OrtException {
        OrtEnvironment env = OrtEnvironment.getEnvironment();
        OnnxTensor inputTensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(inputVector), new long[]{1, inputVector.length});
        Map<String, OnnxTensor> inputs = Collections.singletonMap("input", inputTensor);

        OrtSession.Result results = session.run(inputs);
        float[] output = (float[]) results.get(0).getValue();

        return output[0]; // 返回 AI 评分（0~1），比如：0.8 代表是垃圾内容，0.2 代表是合规文章
    }

    public static void main(String[] args) throws Exception {
        // 1. 初始化 ONNX 模型
        ONNXTextClassifier classifier = new ONNXTextClassifier("bert-classifier.onnx");

        // 2. 假设输入是一个文本特征向量（需要预处理）
        float[] textVector = {0.1f, 0.3f, 0.5f, 0.7f};  // 示例向量，真实情况下需用 BERT 处理文本

        // 3. 进行分类
        float score = classifier.classifyText(textVector);
        System.out.println("AI 评分：" + score);

        // 4. 判断是否违规
        if (score > 0.7) {
            System.out.println("文章违规，需人工审核！");
        } else {
            System.out.println("文章合规，可以发布！");
        }
    }
}
